
from pathlib import Path
import sys
import matplotlib.pyplot as plt
import wave
import numpy as np

py_root = Path(__file__).parent.parent.resolve()
sys.path.append(str(py_root))
# %% test
if __name__ == "__main__":
    from afsk import AudioCPFSK, signal2indicator, indicator2bytes
    # 调制参数
    f_sample = 16_000      # 采样频率
    f_space = 1500         # 空间频率 (0)
    f_mark = 1800          # 标记频率 (1)
    bit_rate = 200         # 比特率
    
    # 创建CPFSK调制器
    modulator = AudioCPFSK(f_sample, f_space, f_mark, bit_rate)
    
    byte_seq = '\x01\x02hello\x10world\x03\x04'.encode('utf-8')
    test_bits = ''.join(format(byte, '08b') for byte in byte_seq)
    test_bits = [bit == '1' for bit in test_bits]  # 转bool
    
    # 进行调制
    attack_bits = 16  # 从调制器调用中获取
    modulated_signal = modulator(test_bits, 0.4, attack_bits=attack_bits)
    
    # 生成时间轴
    duration = len(modulated_signal) / f_sample
    time_axis = np.linspace(0, duration, len(modulated_signal))
        
    # 解调参数
    S_GOERTZEL = 9 # Goertzel算法的窗口系数
    t_per_bit = 1 / bit_rate  # Duration per bit (位持续时间)
    
    # 解调信号(Goertzel算法)
    signal_f1, signal_f0, indicators = signal2indicator(
        modulated_signal,
        fs=f_sample,
        f0=f_space,
        f1=f_mark,
        duration=t_per_bit,
        s_goertzel=S_GOERTZEL
    )

    # 将指示器转换为字节
    results = indicator2bytes(
        indicators,
        start=b'\x01\x02',
        end=b'\x03\x04',
        seperate=b'\x10',
        threshold=0.5
    )

    # 打印解调结果
    print("\n解调结果:")
    for i, r in enumerate(results):
        try:
            decoded_text = r.decode('utf-8')
            print(f"结果 {i}: '{decoded_text}' (原始字节: {r.hex()})")
        except UnicodeDecodeError:
            print(f"结果 {i}: 无法解码为UTF-8 (原始字节: {r.hex()})")
        
    # %% 创建图形
    # 调制部分
    plt.figure(figsize=(12, 8))
    
    # 计算时间偏移量（前导载波持续时间）
    preamble_duration = attack_bits / bit_rate
    data_duration = len(test_bits) / bit_rate
    
    # 子图1: 原始比特序列（数据部分，前导用虚线表示）
    plt.subplot(3, 1, 1)
    
    # 启动信号部分 - 显示为中心频率标识
    plt.axvspan(0, preamble_duration, alpha=0.3, color='orange', label='启动信号(中心频率)')
    
    # 数据比特部分 - 绘制实际比特序列
    for i, bit in enumerate(test_bits):
        start_time = preamble_duration + i / bit_rate
        end_time = preamble_duration + (i + 1) / bit_rate
        bit_time = np.linspace(start_time, end_time, 10)
        bit_value = np.full(10, int(bit))
        plt.plot(bit_time, bit_value, 'b-', linewidth=2)
        # 在每个比特中心添加数值标签
        plt.text((start_time + end_time)/2, int(bit) + 0.1, str(int(bit)), 
                ha='center', va='bottom', fontsize=8)
    
    plt.ylabel('比特值')
    plt.title('CPFSK调制示例')
    plt.grid(True, alpha=0.3)
    plt.ylim(-0.5, 1.5)
    plt.xlim(0, duration)
    plt.legend()
    
    # 子图2: 调制信号
    plt.subplot(3, 1, 2)
    plt.plot(time_axis, modulated_signal, 'r-', linewidth=1)
    plt.ylabel('幅度')
    plt.xlabel('时间 (秒)')
    plt.title('CPFSK调制信号')
    
    # 标记前导载波区域
    plt.axvspan(0, preamble_duration, alpha=0.2, color='orange', label='前导载波')
    
    # 数据区域背景着色
    bit_1_labeled = False
    bit_0_labeled = False
    for i, bit in enumerate(test_bits):
        start_time = preamble_duration + i / bit_rate
        end_time = preamble_duration + (i + 1) / bit_rate
        if bit:
            label = '比特1' if not bit_1_labeled else ""
            plt.axvspan(start_time, end_time, alpha=0.2, color='blue', label=label)
            bit_1_labeled = True
        else:
            label = '比特0' if not bit_0_labeled else ""
            plt.axvspan(start_time, end_time, alpha=0.2, color='gray', label=label)
            bit_0_labeled = True
    
    plt.grid(True, alpha=0.3)
    plt.xlim(0, duration)
    plt.legend()
    
    # 子图3: 频谱分析
    plt.subplot(3, 1, 3)
    frequencies = np.fft.fftfreq(len(modulated_signal), 1/f_sample)
    fft_signal = np.abs(np.fft.fft(modulated_signal))
    # 只显示正频率部分
    positive_freq_idx = frequencies >= 0
    plt.plot(frequencies[positive_freq_idx], fft_signal[positive_freq_idx])
    plt.ylabel('幅度')
    plt.xlabel('频率 (Hz)')
    plt.title('频谱')
    plt.grid(True, alpha=0.3)
    plt.xlim(0, 4000)  # 显示0-4kHz范围
    
    # 标记关键频率
    plt.axvline(x=f_space, color='g', linestyle='--', label=f'Space频率: {f_space}Hz')
    plt.axvline(x=f_mark, color='m', linestyle='--', label=f'Mark频率: {f_mark}Hz')
    plt.axvline(x=(f_space + f_mark)/2, color='orange', linestyle='--', label=f'中心频率: {(f_space + f_mark)/2}Hz')
    plt.legend()
    
    plt.tight_layout()
    plt.show(block=False)
    
    # 打印信息
    print(f"前导载波: {attack_bits} 比特, 持续时间: {preamble_duration:.3f} 秒")
    print(f"测试比特序列: {''.join(['1' if b else '0' for b in test_bits]):<10}")
    print(f"信号长度: {len(modulated_signal)} 采样点")
    print(f"信号持续时间: {duration:.3f} 秒")
    print(f"Space频率: {f_space} Hz, Mark频率: {f_mark} Hz, 中心频率: {(f_space + f_mark)/2} Hz")
    
    # 解调部分
    # 生成时间轴
    duration_total = len(modulated_signal) / f_sample
    time_signal = np.linspace(0, duration_total, len(modulated_signal))
    time_goertzel = np.linspace(0, duration_total, len(signal_f1))
    time_indicators = np.linspace(0, duration_total, len(indicators))

    # %% 创建综合图形
    plt.figure(figsize=(12, 6))
    
    # 子图1: 原始信号
    plt.subplot(3, 1, 1)
    plt.plot(time_signal, modulated_signal, 'b-', linewidth=0.5, alpha=0.7)
    plt.ylabel('幅度')
    plt.title('原始信号')
    plt.grid(True, alpha=0.3)
    plt.xlim(0, duration_total)
    
    # 子图2: Goertzel幅度检测
    plt.subplot(3, 1, 2)
    plt.plot(time_goertzel, signal_f0, 'g-', linewidth=1, label=f'F0({f_space}Hz) 幅度')
    plt.plot(time_goertzel, signal_f1, 'r-', linewidth=1, label=f'F1({f_mark}Hz) 幅度')
    plt.ylabel('Goertzel幅度')
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.xlim(0, duration_total)

    # 子图3: 频谱分析
    plt.subplot(3, 1, 3)
    frequencies = np.fft.fftfreq(len(modulated_signal), 1/f_sample)
    fft_signal = np.abs(np.fft.fft(modulated_signal))
    positive_freq_idx = frequencies >= 0
    plt.plot(frequencies[positive_freq_idx], fft_signal[positive_freq_idx], 'orange', linewidth=2)
    plt.axvline(x=f_space, color='g', linestyle='--', linewidth=1, label=f'Space频率: {f_space}Hz')
    plt.axvline(x=f_mark, color='r', linestyle='--', linewidth=1, label=f'Mark频率: {f_mark}Hz')
    plt.ylabel('频谱幅度')
    plt.xlabel('频率 (Hz)')
    plt.title('频谱分析')
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.xlim(0, 3000)
    
    plt.tight_layout()
    plt.show(block=False)
    
    # 如果有解调结果，创建比特序列图
    if indicators:
        plt.figure(figsize=(12, 6))
        
        # 转换为比特序列
        bit_sequence = ''.join(['1' if p > 0.5 else '0' for p in indicators])
        
        plt.subplot(2, 1, 1)
        plt.plot(time_indicators, indicators, 'purple', linewidth=2)
        plt.axhline(y=0.5, color='red', linestyle='--', alpha=0.7)
        
        # 背景着色显示比特值
        for i, (t, p) in enumerate(zip(time_indicators, indicators)):
            bit = 1 if p > 0.5 else 0
            color = 'lightblue' if bit == 1 else 'lightgray'
            if i < len(time_indicators) - 1:
                plt.axvspan(t, time_indicators[i+1], alpha=0.3, color=color)
        
        plt.ylabel('P1概率')
        plt.title('比特判决过程')
        plt.grid(True, alpha=0.3)
        plt.ylim(0, 1)
        
        plt.subplot(2, 1, 2)
        # 显示比特序列（前N位，避免过长）
        display_n = 200
        display_bits = bit_sequence[:display_n] if len(bit_sequence) > display_n else bit_sequence
        bit_values = [int(b) for b in display_bits]
        bit_times = np.linspace(0, len(display_bits)/len(bit_sequence) * duration_total, len(display_bits))
        
        plt.step(bit_times, bit_values, 'b-', linewidth=2, where='post')
        plt.ylabel('比特值')
        plt.xlabel('时间 (秒)')
        plt.title(f'解调比特序列 (前{len(display_bits)}位)')
        plt.grid(True, alpha=0.3)
        plt.ylim(-0.5, 1.5)
        
        plt.tight_layout()
        plt.show(block=False)
    
    # 打印统计信息
    print(f"信号长度: {len(modulated_signal)} 采样点")
    print(f"信号持续时间: {duration_total:.3f} 秒")
    print(f"解调出的比特数: {len(indicators)}")
    print(f"平均比特率: {len(indicators)/duration_total:.1f} bps")